package com.test.springai.service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

/**
 * Created with IntelliJ IDEA.
 *
 * @description: 在传统的 RAG 系统中，我们通常直接使用用户的原始问题（或经过简单改写的问题）去向量数据库中进行一次检索。然而，用户的问题可能存在以下问题：
 *     <p>歧义性： 一个问题可能有多种解释。 复杂性： 一个问题可能包含多个子问题或需要从不同角度理解。 关键词缺失： 用户提问的措辞可能与知识库中的文档使用的词汇不完全匹配。 不完整性：
 *     原始问题可能没有提供足够的信息来找到最相关的文档。 所以我们需要将用户的输入分解成多个查询，从而提高检索的准确性和召回率。
 * @author: liuziyang
 * @since: 9/3/25 3:02 PM
 * @modifiedBy:
 * @version: 1.0
 */
@Slf4j
@Service
public class AdvancedRagService {
  private final ChatClient chatClient;
  private final MultiQueryExpander multiQueryExpander;
  private final QueryRewriter queryRewriter;
  private final TranslationQueryTransformer translationQueryTransformer;
  // 最终回答的通用 RAG 提示模板
  private final String FINAL_ANSWER_PROMPT =
      """
            请根据提供的上下文信息回答问题。
            如果上下文没有包含足够的信息来回答问题，请说明你无法回答，不要编造信息。
            上下文:
            {context}
            用户问题: {user_input}
            """;

  /**
   * 构造函数，用于初始化 RagService。 注入 ChatClient.Builder、MultiQueryExpander、QueryRewriter 和
   * TranslationQueryTransformer。
   *
   * @param chatClientBuilder 用于构建 ChatClient 实例的 ChatClient.Builder
   * @param multiQueryExpander 多查询扩展器实例
   * @param queryRewriter 查询重写器实例
   * @param translationQueryTransformer 查询翻译器实例
   */
  @Autowired
  public AdvancedRagService(
      ChatClient.Builder chatClientBuilder,
      MultiQueryExpander multiQueryExpander,
      QueryRewriter queryRewriter,
      TranslationQueryTransformer translationQueryTransformer) {
    this.chatClient = chatClientBuilder.build();
    this.multiQueryExpander = multiQueryExpander;
    this.queryRewriter = queryRewriter;
    this.translationQueryTransformer = translationQueryTransformer;
  }

  /**
   * 执行高级 RAG 流程，包括翻译、查询重写、多查询扩展和文档检索，最终生成 LLM 响应。
   *
   * @param userInput 用户的原始输入问题。
   * @return LLM 基于处理后的查询和检索到的上下文生成的回答。
   */
  public Flux<String> advancedGenerateResponse(String userInput) {
    log.info("Starting advanced RAG process for user input: {}", userInput);
    // 1. 翻译用户输入
    // 假设 translateQuery 方法的参数是 sourceLanguage 和 targetLanguage，这里以 "en" 到 "zh" 为例
    // 实际应用中，sourceLanguage 和 targetLanguage 应根据业务逻辑或用户设置动态确定
    String userInputTranslated = processQueryTranslate(userInput, "en", "zh");
    log.info("User input translated to: {}", userInputTranslated);
    // 2. 重写翻译后的用户输入
    String questionRewritten = processQueryRewrite(userInputTranslated);
    log.info("Query rewritten to: {}", questionRewritten);
    // 3. 将重写后的问题转换为子查询，并从向量数据库中检索相关文档，生成上下文
    String context = processQuestionToSubQuery(questionRewritten);
    log.info("Context generated with length: {}", context.length());
    // 4. 构建最终 Prompt 并调用 LLM
    return getFinalAnswerFromLLM(userInput, context);
  }

  /**
   * 将用户问题转换为子查询，并从向量数据库中检索相关文档。 此方法利用 MultiQueryExpander 实现多查询生成和并行检索。
   *
   * @param question 经过重写或翻译后的用户问题。
   * @return 包含检索到的文档内容的上下文字符串，如果未找到相关信息，则返回提示信息。
   */
  private String processQuestionToSubQuery(String question) {
    log.info("Processing question to sub-query and retrieve documents for: {}", question);
    // 1. 使用 MultiQueryExpander 异步生成子查询
    CompletableFuture<List<String>> subQueriesFuture =
        multiQueryExpander.generateSubQueries(question);
    // 2. 等待子查询生成完毕，然后使用 MultiQueryExpander 异步并行检索文档
    // thenCompose 用于在 CompletableFuture 完成后，使用其结果作为下一个 CompletableFuture 的输入
    CompletableFuture<List<Document>> relevantDocumentsFuture =
        subQueriesFuture.thenCompose(
            // 链式调用 retrieveDocuments 方法
            multiQueryExpander::retrieveDocuments);
    // 3. 阻塞等待所有检索完成，然后进行文档去重
    // .join() 会阻塞当前线程直到 CompletableFuture 完成并返回结果
    List<Document> allRelevantDocuments = relevantDocumentsFuture.join();
    List<Document> uniqueDocuments = multiQueryExpander.deduplicateDocuments(allRelevantDocuments);
    log.info("Final unique documents for context: {}", uniqueDocuments.size());
    // 4. 将去重后的文档内容拼接成上下文
    // 使用双换行符和“---”作为分隔符，提高不同文档在上下文中的可读性
    String context =
        uniqueDocuments.stream()
            // 从 Document 对象中获取文本内容
            .map(Document::getText)
            // 用特殊分隔符区分不同文档
            .collect(Collectors.joining("\n\n---\n\n"));
    // 如果没有检索到相关文档，返回一个默认提示
    if (context.trim().isEmpty()) {
      log.warn("No relevant context found for question: {}", question);
      return "抱歉，我未能找到相关的上下文信息来回答您的问题。";
    }
    return context;
  }

  /**
   * 将用户原始问题和聚合的上下文发送给 LLM 获取最终答案。
   *
   * @param userInput 用户的原始问题（通常用于最终向用户展示或作为 LLM 提示的一部分）。
   * @param context 聚合后的上下文信息。
   * @return LLM 的最终回答。
   */
  private Flux<String> getFinalAnswerFromLLM(String userInput, String context) {
    log.info("Getting final answer from LLM for user input: {}", userInput);
    SystemPromptTemplate finalAnswerTemplate = new SystemPromptTemplate(FINAL_ANSWER_PROMPT);
    // 填充模板中的 {context} 和 {user_input} 占位符
    String systemMessageContent =
        finalAnswerTemplate
            .createMessage(
                Map.of(
                    "context", context,
                    "user_input", userInput))
            .getText();
    Prompt prompt =
        new Prompt(
            List.of(
                // 构建包含上下文和用户问题的系统消息
                new SystemMessage(systemMessageContent)));
    return chatClient.prompt(prompt).stream().content();
  }

  /**
   * 将用户输入进行重写。 此方法利用 QueryRewriter 服务将查询改写为更适合检索的格式。
   *
   * @param userInput 待重写的用户输入。
   * @return 重写后的查询字符串。
   */
  private String processQueryRewrite(String userInput) {
    log.info("Processing query rewrite for: {}", userInput);
    // 使用 QueryRewriter 异步重写查询，并阻塞等待结果
    String rewrittenQuery = queryRewriter.rewriteSingleQuery(userInput).join();
    log.info("Query rewritten to: {}", rewrittenQuery);
    return rewrittenQuery;
  }

  /**
   * 将用户输入进行翻译。 此方法利用 TranslationQueryTransformer 服务将查询从源语言翻译到目标语言。
   *
   * @param userInput 待翻译的用户输入。
   * @param sourceLanguage 源语言（例如 "zh"）。
   * @param targetLanguage 目标语言（例如 "en"）。
   * @return 翻译后的查询字符串。
   */
  private String processQueryTranslate(
      String userInput, String sourceLanguage, String targetLanguage) {
    log.info(
        "Processing query translation for: {} from {} to {}",
        userInput,
        sourceLanguage,
        targetLanguage);
    // 使用 TranslationQueryTransformer 异步翻译查询，并阻塞等待结果
    String translatedQuery =
        translationQueryTransformer
            .translateQuery(userInput, sourceLanguage, targetLanguage)
            .join();
    log.info("Query translated to: {}", translatedQuery);
    return translatedQuery;
  }
}
